package typ

import (
	"errors"
	"fmt"
)

// Environment represents the Scheme environment, storing variable and function bindings
type Environment struct {
	outer *Environment      // outer environment (parent scope)
	store map[string]Object // current environment's variable bindings
}

// NewEnvironment creates a new environment
func NewEnvironment() *Environment {
	env := &Environment{
		outer: nil,
		store: make(map[string]Object),
	}
	return env
}

// NewEnclosedEnvironment creates a new nested environment
func NewEnclosedEnvironment(outer *Environment) *Environment {
	env := &Environment{
		outer: outer,
		store: make(map[string]Object),
	}
	return env
}

// Get looks up a variable in the environment
func (e *Environment) Get(name string) (Object, error) {
	val, ok := e.store[name]
	if !ok && e.outer != nil {
		// Look up in the outer environment
		return e.outer.Get(name)
	}
	if !ok {
		return nil, errors.New(fmt.Sprintf("unbound variable: %s", name))
	}
	return val, nil
}

// Set sets a variable in the current environment
func (e *Environment) Set(name string, val Object) Object {
	e.store[name] = val
	return val
}

// DefineBuiltins defines built-in functions
func DefineBuiltins(env *Environment) {
	// Mathematical operations
	env.Set("+", &Builtin{Fn: builtinAdd})
	env.Set("-", &Builtin{Fn: builtinSubtract})
	env.Set("*", &Builtin{Fn: builtinMultiply})
	env.Set("/", &Builtin{Fn: builtinDivide})

	// Comparison operations
	env.Set(">", &Builtin{Fn: builtinGreater})
	env.Set("<", &Builtin{Fn: builtinLess})
	env.Set("=", &Builtin{Fn: builtinEqual})
	env.Set(">=", &Builtin{Fn: builtinGreaterEqual})
	env.Set("<=", &Builtin{Fn: builtinLessEqual})

	// Logical operations
	env.Set("and", &Builtin{Fn: builtinAnd})
	env.Set("or", &Builtin{Fn: builtinOr})
	env.Set("not", &Builtin{Fn: builtinNot})

	// List operations
	env.Set("cons", &Builtin{Fn: builtinCons})
	env.Set("car", &Builtin{Fn: builtinCar})
	env.Set("cdr", &Builtin{Fn: builtinCdr})
	env.Set("list", &Builtin{Fn: builtinList})
	env.Set("null?", &Builtin{Fn: builtinIsNull})
	env.Set("eq?", &Builtin{Fn: builtinEq})
	env.Set("equal?", &Builtin{Fn: builtinEqual})

	// Other built-in functions
	env.Set("display", &Builtin{Fn: builtinDisplay})
	env.Set("newline", &Builtin{Fn: builtinNewline})
}

// Mathematical operation built-in functions
func builtinAdd(args []Object) Object {
	// Check if all arguments are numbers
	for _, arg := range args {
		if !isNumber(arg) {
			return &String{Value: "Error: + expects numbers"}
		}
	}
	
	// Check if all arguments are Rational
	allRational := true
	for _, arg := range args {
		if _, ok := arg.(*Rational); !ok {
			allRational = false
			break
		}
	}
	
	// If all are Rational, perform exact rational arithmetic
	if allRational && len(args) > 0 {
		result := args[0].(*Rational)
		for i := 1; i < len(args); i++ {
			result = result.Add(args[i].(*Rational))
		}
		return result
	}
	
	// Otherwise, use floating point arithmetic
	var result float64 = 0
	for _, arg := range args {
		result += arg.(Number).Value()
	}
	return &Real{RealValue: result}
}

func builtinSubtract(args []Object) Object {
	if len(args) == 0 {
		return &String{Value: "Error: - requires at least one argument"}
	}
	// Check if the first argument is a number
	if !isNumber(args[0]) {
		return &String{Value: "Error: - expects numbers"}
	}
	
	// If there's only one argument, return its negative value
	if len(args) == 1 {
		if rational, ok := args[0].(*Rational); ok {
			return &Rational{Numerator: -rational.Numerator, Denominator: rational.Denominator}
		}
		return &Real{RealValue: -args[0].(Number).Value()}
	}
	
	// Check if all arguments are Rational
	allRational := true
	for _, arg := range args {
		if _, ok := arg.(*Rational); !ok {
			allRational = false
			break
		}
	}
	
	// If all are Rational, perform exact rational arithmetic
	if allRational {
		result := args[0].(*Rational)
		for i := 1; i < len(args); i++ {
			result = result.Subtract(args[i].(*Rational))
		}
		return result
	}
	
	// Otherwise, use floating point arithmetic
	// Check remaining arguments
	for i := 1; i < len(args); i++ {
		if !isNumber(args[i]) {
			return &String{Value: "Error: - expects numbers"}
		}
	}
	
	result := args[0].(Number).Value()
	for i := 1; i < len(args); i++ {
		result -= args[i].(Number).Value()
	}
	return &Real{RealValue: result}
}

func builtinMultiply(args []Object) Object {
	// Check if all arguments are numbers
	for _, arg := range args {
		if !isNumber(arg) {
			return &String{Value: "Error: * expects numbers"}
		}
	}
	
	// Check if all arguments are Rational
	allRational := true
	for _, arg := range args {
		if _, ok := arg.(*Rational); !ok {
			allRational = false
			break
		}
	}
	
	// If all are Rational, perform exact rational arithmetic
	if allRational && len(args) > 0 {
		result := args[0].(*Rational)
		for i := 1; i < len(args); i++ {
			result = result.Multiply(args[i].(*Rational))
		}
		return result
	}
	
	// Otherwise, use floating point arithmetic
	var result float64 = 1
	for _, arg := range args {
		result *= arg.(Number).Value()
	}
	return &Real{RealValue: result}
}

func builtinDivide(args []Object) Object {
	if len(args) == 0 {
		return &String{Value: "Error: / requires at least one argument"}
	}
	// Check if the first argument is a number
	if !isNumber(args[0]) {
		return &String{Value: "Error: / expects numbers"}
	}
	
	// If there's only one argument, return 1 divided by that number
	if len(args) == 1 {
		if rational, ok := args[0].(*Rational); ok {
			r1 := &Rational{Numerator: 1, Denominator: rational.Numerator}
			r2 := &Rational{Numerator: rational.Denominator, Denominator: 1}
			return r1.Multiply(r2)
		} else if integer, ok := args[0].(*Integer); ok {
			// For integer, return 1 divided by it as rational
			return &Rational{Numerator: 1, Denominator: integer.IntValue}
		}
		return &Real{RealValue: 1 / args[0].(Number).Value()}
	}
	
	// Check if all arguments are Integer or Rational
	allIntegerOrRational := true
	for _, arg := range args {
		if _, ok := arg.(*Integer); ok {
			continue
		}
		if _, ok := arg.(*Rational); ok {
			continue
		}
		allIntegerOrRational = false
		break
	}
	
	// If all are Integer or Rational, perform exact rational arithmetic
	if allIntegerOrRational {
		// Convert all arguments to Rational
		rationalArgs := make([]*Rational, len(args))
		for i, arg := range args {
			if integer, ok := arg.(*Integer); ok {
				rationalArgs[i] = &Rational{Numerator: integer.IntValue, Denominator: 1}
			} else if rational, ok := arg.(*Rational); ok {
				rationalArgs[i] = rational
			}
		}
		
		// Perform division
		result := rationalArgs[0]
		for i := 1; i < len(rationalArgs); i++ {
			// Check for division by zero
			if rationalArgs[i].Numerator == 0 {
				panic("Error: division by zero")
			}
			result = result.Divide(rationalArgs[i])
		}
		return result
	}
	
	// Otherwise, use floating point arithmetic
	// Check remaining arguments
	for i := 1; i < len(args); i++ {
		if !isNumber(args[i]) {
			return &String{Value: "Error: / expects numbers"}
		}
		divisor := args[i].(Number).Value()
		if divisor == 0 {
			panic("Error: division by zero")
		}
	}
	
	result := args[0].(Number).Value()
	for i := 1; i < len(args); i++ {
		divisor := args[i].(Number).Value()
		result /= divisor
	}
	return &Real{RealValue: result}
}

// isNumber checks if an object is a number type
func isNumber(obj Object) bool {
	_, ok := obj.(Number)
	return ok
}

// Comparison operation built-in functions
func builtinGreater(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: > requires exactly two arguments"}
	}
	if !isNumber(args[0]) || !isNumber(args[1]) {
		return &String{Value: "Error: > expects numbers"}
	}
	a := args[0].(Number).Value()
	b := args[1].(Number).Value()
	if a > b {
		return True
	}
	return False
}

func builtinLess(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: < requires exactly two arguments"}
	}
	if !isNumber(args[0]) || !isNumber(args[1]) {
		return &String{Value: "Error: < expects numbers"}
	}
	a := args[0].(Number).Value()
	b := args[1].(Number).Value()
	if a < b {
		return True
	}
	return False
}

func builtinEqual(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: = requires exactly two arguments"}
	}

	a, b := args[0], args[1]
	
	// Special handling for numbers: compare values regardless of type
	if isNumber(a) && isNumber(b) {
		return booleanObject(a.(Number).Value() == b.(Number).Value())
	}

	// For other types, check if they have the same type first
	if a.Type() != b.Type() {
		return False
	}

	switch a.Type() {
	case STRING:
		return booleanObject(a.(*String).Value == b.(*String).Value)
	case BOOLEAN:
		return booleanObject(a.(*Boolean).Value == b.(*Boolean).Value)
	case SYMBOL:
		return booleanObject(a.(*Symbol).Value == b.(*Symbol).Value)
	case NULL:
		return True
	case PAIR:
		// Simple equality check, for lists we only check reference equality
		return booleanObject(a == b)
	default:
		return False
	}
}

func builtinGreaterEqual(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: >= requires exactly two arguments"}
	}
	if !isNumber(args[0]) || !isNumber(args[1]) {
		return &String{Value: "Error: >= expects numbers"}
	}
	a := args[0].(Number).Value()
	b := args[1].(Number).Value()
	if a >= b {
		return True
	}
	return False
}

func builtinLessEqual(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: <= requires exactly two arguments"}
	}
	if !isNumber(args[0]) || !isNumber(args[1]) {
		return &String{Value: "Error: <= expects numbers"}
	}
	a := args[0].(Number).Value()
	b := args[1].(Number).Value()
	if a <= b {
		return True
	}
	return False
}

// Logical operation built-in functions
func builtinAnd(args []Object) Object {
	for _, arg := range args {
		if isFalse(arg) {
			return False
		}
	}
	if len(args) > 0 {
		return args[len(args)-1]
	}
	return True
}

func builtinOr(args []Object) Object {
	for _, arg := range args {
		if !isFalse(arg) {
			return arg
		}
	}
	return False
}

func builtinNot(args []Object) Object {
	if len(args) != 1 {
		return &String{Value: "Error: not requires exactly one argument"}
	}
	if isFalse(args[0]) {
		return True
	}
	return False
}

// List operation built-in functions
func builtinCons(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: cons requires exactly two arguments"}
	}
	return &Pair{Car: args[0], Cdr: args[1]}
}

func builtinCar(args []Object) Object {
	if len(args) != 1 {
		return &String{Value: "Error: car requires exactly one argument"}
	}
	if args[0].Type() != PAIR {
		return &String{Value: "Error: car expects a pair"}
	}
	return args[0].(*Pair).Car
}

func builtinCdr(args []Object) Object {
	if len(args) != 1 {
		return &String{Value: "Error: cdr requires exactly one argument"}
	}
	if args[0].Type() != PAIR {
		return &String{Value: "Error: cdr expects a pair"}
	}
	return args[0].(*Pair).Cdr
}

func builtinList(args []Object) Object {
	var result Object = Nil
	for i := len(args) - 1; i >= 0; i-- {
		result = &Pair{Car: args[i], Cdr: result}
	}
	return result
}

func builtinIsNull(args []Object) Object {
	if len(args) != 1 {
		return &String{Value: "Error: null? requires exactly one argument"}
	}
	return booleanObject(args[0].Type() == NULL)
}

func builtinEq(args []Object) Object {
	if len(args) != 2 {
		return &String{Value: "Error: eq? requires exactly two arguments"}
	}
	// eq? checks reference equality in Scheme
	return booleanObject(args[0] == args[1])
}

// Other helper functions
func builtinDisplay(args []Object) Object {
	for _, arg := range args {
		fmt.Print(arg.String())
	}
	return Nil
}

func builtinNewline(args []Object) Object {
	fmt.Println()
	return Nil
}

// Helper functions
func booleanObject(b bool) Object {
	if b {
		return True
	}
	return False
}

func isFalse(obj Object) bool {
	if obj == False || obj.Type() == NULL {
		return true
	}
	return false
}
